import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
import os
from PIL import Image
import matplotlib.pyplot as plt
from .utils import Transform_Two_RGB, de_normalizer


class RainTrainH(Dataset):

    def __init__(self, root=r"E:\note\ssl\data\ACDC", split="train", transform=None):
        super(RainTrainH, self).__init__()
        self.split = split
        self.root = root
        self.transform = transform
        self.sample_list = []
        self.load_annotations()  # 加载文件路径
        print("total {} samples".format(len(self.sample_list)))

    def __len__(self):
        return len(self.sample_list)

    def __getitem__(self, idx):
        case = self.sample_list[idx]
        image = Image.open(os.path.join(self.root, self.split, "image", case))
        mask = Image.open(os.path.join(self.root, self.split, "target", case))

        if self.transform is not None:
            result = self.transform(image=image, mask=mask)
            image = result["image"]
            mask = result["mask"]

        return image, mask

    def load_annotations(self):
        self.sample_list = np.array(os.listdir(os.path.join(self.root, self.split, "image")))


def get_raintrainh_loader(root=r'/home/ubuntu/data/ACDC', batch_size=4, train_crop_size=(224, 224)):
    """

    :param root:
    :param batch_size: 批次大小
    :param label: 有标签的数量
    :return:
    """
    train_transform = Transform_Two_RGB(mode="train", out_size=train_crop_size)
    test_transform = Transform_Two_RGB(mode="test", out_size=train_crop_size)
    train_dataset = RainTrainH(root=root, split="train", transform=train_transform)
    test_dataset = RainTrainH(root=root, split="test", transform=test_transform)
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, num_workers=4, shuffle=True, drop_last=True)
    test_dataloader = DataLoader(test_dataset, batch_size=1, num_workers=4, shuffle=False)
    return train_dataloader, test_dataloader


def show(im):
    # im = torch.clamp(im, 0.0, 1.0)
    # im = im.permute(1, 2, 0).numpy()
    im = de_normalizer(im)
    fig = plt.figure()

    # print(im)
    plt.imshow(im)
    plt.show()

    # skimage.io.imsave("1.png",im)
    # fig.savefig("result.png")


def show_label(mask, path="label.jpg"):
    plt.figure()
    plt.imshow(mask)
    plt.show()
    Image.fromarray(mask).save(path)


# def get_ssl_acdc_loader(root=r'/home/ubuntu/data/ACDC', batch_size=8, unlabel_batch_size=24, train_crop_size=(224, 224), label_num=0.2,seed=1337):
#     """
#     :param root: 数据集路径
#     :param batch_size: 有标注数据批次大小
#     :param unlabel_batch_size: 无标注数据的batch大小
#     :param label_num: 有标签的数量
#     :return:
#     """

#     def worker_init_fn(worker_id):
#         random.seed(seed + worker_id)

#     train_transform=RandomGenerator(train_crop_size)
#     train_dataset = ACDC(root=root, split="train", transform=train_transform)
#     total_slices=len(train_dataset)
#     labeled_slice=patients_to_slices("ACDC",label_num)
#     labeled_idxs = list(range(0, labeled_slice))
#     unlabeled_idxs = list(range(labeled_slice, total_slices))

#     batch_sampler = TwoStreamBatchSampler(labeled_idxs, unlabeled_idxs, batch_size+unlabel_batch_size,unlabel_batch_size)
#     train_loader=DataLoader(train_dataset,batch_sampler=batch_sampler,num_workers=4,pin_memory=True,worker_init_fn=worker_init_fn)

#     test_dataset = ACDC(root=root, split="test")
#     test_loader = DataLoader(test_dataset, batch_size=1, num_workers=4, shuffle=False)
#     return train_loader, test_loader


if __name__ == '__main__':

    train_dataloader, test_dataloader = get_raintrainh_loader(root=r"D:\BaiduNetdiskDownload\datasets\RainTrainH-My")
    # print(len(train_dataloader))
    # print(len(test_dataloader))
    # print(len(test_dataloader.dataset))
    for image, label in train_dataloader:
        print(image.shape)
        print(label.shape)
        # print(np.unique(label.numpy()))
        show(image[0])
        show(label[0])
        # show_label(train_dataloader.dataset.label_to_img(label))
        break

    for sample in test_dataloader:
        image, label = sample
        print(image.shape)
        print(label.shape)
        show(image[0])
        show(label[0])
        # print(np.unique(label.numpy()))
        # show(image[0])
        # show_label(label[0].numpy())
        break
